package ClientPackage;

import SecureUtil.DiffieHellman;
import SecureUtil.PBEHelper;
import SecureUtil.Packer;
import SecureUtil.RSAHelper;
import SecureUtil.UDPHelper;
import SecureUtil.Packer.DataType;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.net.*;
import java.security.AlgorithmParameterGenerator;
import java.security.AlgorithmParameters;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.InvalidParameterSpecException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;

import javax.crypto.SecretKey;
import javax.crypto.spec.DHParameterSpec;

public class Client {

	private static Client _singleton = null;
	// Server
	private InetAddress _serverIp;
	private int _serverPort;
	private PublicKey _serverpubKey;
	private InetAddress _serverLoginIp;
	private int _serverLoginPort;

	private PrivateKey _clientPrivKey;
	private PublicKey _clientPubKey;
	private SecretKey _serverSessionKey;

	// identity information
	private String _userName;
	private String _token;
	private RSAHelper _clientRSA;
	// local socket object including local ip and port;
	private DatagramSocket _clientSocket;

	// online user
	private ArrayList<ClientInfo> _onlineUserList;

	// session information

	private final String LOGIN = "LOGIN";
	private final String LOGOUT = "LOGOUT";
	private final String UPDATE_LIST = "UPDATE_LIST";

	private Client(String configfile) throws Exception {
		this.init(configfile);
		_onlineUserList = new ArrayList<ClientInfo>();
	}

	public static Client getInstance(String configfile) throws Exception {
		if (_singleton == null)
			_singleton = new Client(configfile);
		return _singleton;
	}

	public static Client getInstance() throws Exception {
		if (_singleton == null)
			_singleton = new Client("config.txt");
		return _singleton;
	}

	private void init(String configfile) throws Exception {
		_clientSocket = new DatagramSocket(); // generate a socket random port
		// number and a ip
		InputStream in = new FileInputStream(configfile);
		BufferedReader br = new BufferedReader(new InputStreamReader(in));

		_serverIp = InetAddress.getByName(br.readLine().substring(12));
		System.out.println("get server Ip from config file: " + _serverIp);
		_serverPort = Integer.parseInt(br.readLine().substring(14));
		System.out.println("get server Port from config file: " + _serverPort);
		_clientRSA = new RSAHelper();
		_clientRSA.generateKeys(); // generate two key pairs(for Public Key
		_clientPubKey = _clientRSA.pubKey;
		_clientPrivKey = _clientRSA.privKey;

		_serverpubKey = _clientRSA.getPublicKeyFromFile("Server_public_key");

	}

	/**
	 * This method is used for log in
	 */
	private boolean login(String pw) throws Exception {
		int maxtrialNum = 5;
		while (maxtrialNum > 0) {
			byte[] N1 = LoginMessage1234(pw);
			// System.out.println("test2");
			if (N1 == null) {
				System.out.println("login retrying " + (maxtrialNum - 1)
						+ " left");
				maxtrialNum--;
				continue;
			}
			byte[] N2 = new byte[20];
			new SecureRandom().nextBytes(N2);
			if (!LoginMessage56(N1, N2)) {
				System.out.println("login retrying " + (maxtrialNum - 1)
						+ " left");
				maxtrialNum--;
				continue;
			}
			return true;
		}
		return false;
	}

	private boolean logout() {
		try {
			DiffieHellman dh = new DiffieHellman();
			byte[] encryptedLogout = dh.encryptDH(LOGOUT.getBytes(),
					_serverSessionKey);
			byte[] bye = sendCommand(LOGOUT, encryptedLogout);
			return Arrays.equals(bye, LOGOUT.getBytes());
		} catch (SocketException e) {
			e.printStackTrace();
			return false;
		} catch (Exception e) {
			e.printStackTrace();
			return false;
		}
	}

	public synchronized boolean updateList() {
		try {

			DiffieHellman dh = new DiffieHellman();
			// byte[] msg = dh.encryptDH(UPDATE_LIST.getBytes(),
			// _serverSessionKey);
			byte[] cipher = sendCommand(UPDATE_LIST, null);
			if (cipher == null)
				return false;
			byte[] rawList = dh.decryptDH(cipher, _serverSessionKey);
			Packer pack = new Packer();
			pack.unpack(rawList);
			int num = (Integer) pack.getNext(DataType.INTEGER);
			this._onlineUserList.clear();
			for (int i = 0; i < num; i++)
				updateUserInfo(new ClientInfo((byte[]) pack
						.getNext(DataType.BYTES)));

		} catch (SocketException e) {
			e.printStackTrace();
			return false;
		} catch (Exception e) {
			e.printStackTrace();
			return false;
		}

		return true;
	}

	void printOnlineList() {
		for (ClientInfo c : _onlineUserList) {
			System.out.printf("%s\t%s\t%d\t%s\t%s\n", c._username, c._ip
					.getHostAddress(), c._port,
					(c._pubkey != null) ? "w/ pubkey" : "w/o Pubkey",
					(c._sessionkey != null) ? "w/ Sesskey" : "w/o Sesskey");
		}
	}

	private void updateUserInfo(ClientInfo c) {
		_onlineUserList.add(c);
	}

	/**
	 * This method is used to send message to other users
	 * 
	 * @param username
	 * @param message
	 * @return
	 * @throws Exception
	 */
	private boolean SendMessage(String username, String message)
			throws Exception {
		ClientInfo to = SearchClientInfo(username);
		if (to == null) {
			System.out.println("user " + username
					+ " is not online or is not a valid name.");
			return false;
		} else {
			System.out.println("found use in the onlineUserList: " + username
					+ "    ready to send message.");
			boolean sent = secureSendMessage(to, message);
			if (sent) {
				System.out.println("Communication established, Message sent!");
				return true;

			} else {
				System.out.println("Fail to connect, message can not be sent.");
				return false;
			}
		}
	}

	/**
	 * This message implements the send message in the secure protocol
	 * 
	 * @param anotherClient
	 *            This is the clientInfo that you want send message to
	 * @param messagePlainText
	 * @return
	 */
	private boolean secureSendMessage(ClientInfo anotherClient,
			String messagePlainText) {
		try {
			// this socket used to send message to server
			DatagramSocket testSocket = new DatagramSocket();
			// this udphelper is used for connect
			UDPHelper udphelper = new UDPHelper(testSocket);

			// 1) send request without cookie ---1st
			String request = "request";
			String noCookie = "";
			Packer packer = new Packer();
			packer.addElement(DataType.STRING, request);
			packer.addElement(DataType.STRING, noCookie);
			udphelper.sendBytes(packer.pack(), anotherClient._ip,
					anotherClient._port);
			packer.Reset();
			System.out.println("sent request without cookie.");

			// 2) wait for cookie ---2nd
			byte[] receivedCookie = udphelper.receiveBytes(10000);
			System.out.println("received a cookie form client "
					+ anotherClient._username);
			// 3) send request with cookie ---3rd
			String request2 = "request2";
			String cookie = new String(receivedCookie);
			packer.addElement(DataType.STRING, request2);
			packer.addElement(DataType.STRING, cookie);
			udphelper.sendBytes(packer.pack(), udphelper.fromip,
					udphelper.fromport);
			packer.Reset();
			System.out.println("sent request with cookie.");

			// 4) wait for (gb mod p, {N1}A) ---4th
			// At this time, ip and port should be changed to another channel
			byte[] received4th = udphelper.receiveBytes(10000);
			System.out.println("received a (gb mod p, {N1}A) form client "
					+ anotherClient._username);

			// 5) send ga mod p, {A,N2,N1}KB, ---5th

			String username = _userName;

			packer.unpack(received4th);
			byte[] receiveFistHalfDH = (byte[]) packer
					.getNext(Packer.DataType.BYTES);
			byte[] receiveEncryptedN1 = (byte[]) packer
					.getNext(Packer.DataType.BYTES);
			packer.Reset();

			DiffieHellman dh = new DiffieHellman();
			byte[] secondHalfDH = dh.processFirstHalf(receiveFistHalfDH);// DH
																			// second
																			// half

			RSAHelper rsaHelper = new RSAHelper();
			rsaHelper.privKey = this._clientPrivKey;
			byte[] N1 = rsaHelper.RSADecrypt(receiveEncryptedN1);

			SecureRandom random = new SecureRandom();
			byte[] N2 = new byte[24];
			random.nextBytes(N2); // generate N2

			packer.addElement(DataType.STRING, username);
			packer.addElement(DataType.BYTES, N2);
			packer.addElement(DataType.BYTES, N1);

			// {A,N2,N1}KB encrypt with B's public key
			rsaHelper.pubKey = anotherClient._pubkey;
			byte[] encryptedHalf = rsaHelper.RSAEncrypt(packer.pack());
			packer.Reset();

			packer.addElement(DataType.BYTES, secondHalfDH);
			packer.addElement(DataType.BYTES, encryptedHalf);

			udphelper.sendBytes(packer.pack(), udphelper.fromip,
					udphelper.fromport);
			packer.Reset();
			System.out.println("sent (ga mod p, {A,N2,N1}KB) to client "
					+ anotherClient._username);

			// 6) wait for Ka{N2}
			byte[] finalEncryptedBytes = udphelper.receiveBytes(10000);
			System.out.println("received Ka{N2}");

			// 7) send Message!
			SecretKey serverSharedkey = dh.serverSharedKey();
			byte[] decryptedBytes = dh.decryptDH(finalEncryptedBytes,
					serverSharedkey);
			if (!(Arrays.equals(N2, decryptedBytes)))
				throw new Exception("Not the same Diffie-Hellman Key.");
			byte[] encryptedMessage = dh.encryptDH(messagePlainText.getBytes(),
					serverSharedkey);
			udphelper.sendBytes(encryptedMessage, udphelper.fromip,
					udphelper.fromport, serverSharedkey);
			System.out.println("received encrypted message to client "
					+ anotherClient._username);
			return true;

		} catch (IOException e) {

			e.printStackTrace();
			return false;
		} catch (Exception e) {

			e.printStackTrace();
			return false;
		}
	}

	/**
	 * This method is used to find out the certain user is on line or not
	 * 
	 * @param username
	 * @return
	 */
	public ClientInfo SearchClientInfo(String username) {

		// first update the list search again.
		updateList();

		for (ClientInfo c : _onlineUserList) {
			if (c._username.equals(username))
				return c;
		}
		return null;
	}

	// used to listen from network including server's list updating
	// and other clients message request.
	public void run() {

	}

	/**
	 * Main thread is used to listen for user's command First login, you need to
	 * provide user name and password. Once you logged in, you could send
	 * message to server or other on-line users. you need to type "List" first
	 * to get on-line users list so that you know who you could send message to.
	 * 
	 * then it has 3 types of commands: 1.List ----This command is used to list
	 * all the on-line users. 2.Message TA1 hi, how are you? ----This command
	 * will send message to user TA1: hi, how are you? 3.Logout ----This command
	 * is used to logout the user.
	 * 
	 * @param argv
	 */
	public static void main(String[] argv) {
		try {
			Client client = Client.getInstance("config.txt");
			String password = "";
			BufferedReader input = new BufferedReader(new InputStreamReader(
					System.in));

			do {
				System.out.print("Username>>");
				client._userName = input.readLine();
				System.out.print("\nPassword>>");
				password = input.readLine();

			} while (!client.login(password));

			// spawn a new thread for receiving packets
			Thread receiving = new Thread(new ClientReceiveThread(
					client._clientSocket, client._clientPrivKey));
			receiving.start();

			// although clients can try infinit times, server won't give it
			// allow it to login after 3 trials failed in an hour.
			System.out
					.println("Login success! Now you could start chatting with others.");

			/**
			 * Here we loop to wait for commands: List, Message and logout
			 * 1.Send Bob hi, how are you! 2.List 3.Logout
			 */
			Command cmd;

			System.out.println("Input a command(List, Send or Logout): ");
			while ((cmd = new Command()) != null) {

				if (cmd.getCmd().equalsIgnoreCase("List")) {
					System.out.println("On-line users are: ");
					System.out.printf("%s\t%s\t%s\t%s\t%s\n", "User",
							"Ip Address", "port#", "PublicKey", "SessionKey");
					client.updateList();
					client.printOnlineList();
				} else if (cmd.getCmd().equalsIgnoreCase("Send")) {
					System.out.println("get " + cmd.getCmd());
					client.SendMessage(cmd.getUserName(), cmd.getMessage());
				} else if (cmd.getCmd().equalsIgnoreCase("Logout")) {
					System.out.println("get " + cmd.getCmd());
					if (client.logout()) {
						System.out.println("Client " + client._userName
								+ " Logged out.");

						
						System.exit(0);

						//break;
					}

				}
				System.out.println
				("Input a command(List, Send or Logout): ");

			} // end of while
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	/**
	 * @param password
	 *            the original text of password, not hashcode,
	 * @return N1 got from server; session Key shared with Server is established
	 *         ans stored implicitly.
	 * @throws Exception
	 */
	private byte[] LoginMessage1234(String password) throws Exception {
		// generate message3: LOGIN, Cookie, {A}, ga mod p server
		DiffieHellman DHclient = new DiffieHellman();
		byte[] clientPubKey = DHclient.generateFirstHalf();
		Packer packer = new Packer();
		RSAHelper rsaAgent = new RSAHelper();
		rsaAgent.pubKey = _serverpubKey;
		byte[] uNameEncrypted = rsaAgent.RSAEncrypt(_userName.getBytes());

		packer.addElement(Packer.DataType.BYTES, uNameEncrypted);
		packer.addElement(Packer.DataType.BYTES, clientPubKey);

		byte[] initmsg = packer.pack();

		// preocess login message 1,2,3 and get 4
		// get message3: LOGIN, Cookie, {A, ga mod p}server
		byte[] message4 = sendCommand("LOGIN", initmsg);

		if (message4 == null) {
			System.err.println("can not reach server (time out: 10s) "
					+ "try it again");
			return null;
		}
		packer.Reset();

		packer.unpack(message4);
		String salt;
		byte[] WN1DH2;
		try {
			salt = (String) packer.getNext(Packer.DataType.STRING);
			WN1DH2 = (byte[]) packer.getNext(Packer.DataType.BYTES);
		} catch (Exception e) {
			System.err.println("can not unpack login message 4 in client");
			return null;
		}
		// exctract N1 and DH2
		byte[] N1DH2 = PBEHelper.PBEdecrypt(PBEHelper.getPBEKey(PBEHelper
				.getPasswordHash(password), salt), WN1DH2);
		packer.Reset();

		packer.unpack(N1DH2);
		byte[] N1 = null;
		byte[] DH2 = null;
		try {
			N1 = (byte[]) packer.getNext(Packer.DataType.BYTES);
			DH2 = (byte[]) packer.getNext(Packer.DataType.BYTES);
		} catch (Exception e) {
			System.err.println("can not unpack N1,DH2");
			return null;
		}
		_serverSessionKey = DHclient.ClientSharedKey(DH2);
		return N1;
	}

	/**
	 * deal with the message 5, 6;
	 * 
	 * @param N1
	 *            send back to server to authenticat this client
	 * @param N2
	 *            used to authenticate server, when server come back with
	 *            correct N
	 * @return whether server comming back with correct N2
	 * @throws Exception
	 */
	// true if Server comming back with correct N2;
	private boolean LoginMessage56(byte[] N1, byte[] N2) throws Exception {
		// Message5: Ka{A, N1, KA, N2} ; Ka = gab mod p
		Packer packer = new Packer();
		packer.addElement(Packer.DataType.STRING, _userName);
		packer.addElement(Packer.DataType.BYTES, N1);
		packer.addElement(Packer.DataType.BYTES, _clientPubKey.getEncoded());
		packer.addElement(Packer.DataType.BYTES, N2);
		DiffieHellman df = new DiffieHellman();
		byte[] msg5 = df.encryptDH(packer.pack(), _serverSessionKey);

		UDPHelper sendAgent = new UDPHelper(_clientSocket);
		sendAgent.sendBytes(msg5, _serverLoginIp, _serverLoginPort);

		// wait for message 6: Ka{N2};
		byte[] msg6 = sendAgent.receiveBytes(10000);

		byte[] N2recved = df.decryptDH(msg6, _serverSessionKey);
		return Arrays.equals(N2recved, N2);
	}

	/**
	 * This method deals with three message to server: 1. send a request: cmd,
	 * NULL, [EMPTY] 2. get a cookie back: cookie 3. send request with cookie
	 * and arguments: cmd, cookie, argument
	 * 
	 * @param cmd
	 *            the command send to server, it is one of "LOGIN", "LOGOUT",
	 *            "UPDATE_LIST"
	 * @param msg
	 *            the argument for the request; for login, it's init message.
	 *            for LOGOUT, and LIST, it's K{LOGOUT} or K{LIST} for verify the
	 *            request.
	 * @return
	 * @throws Exception
	 */
	private byte[] sendCommand(String cmd, byte[] msg) throws Exception {
		try {
			UDPHelper udphelper;
			if (cmd.equalsIgnoreCase(LOGIN))
				udphelper = new UDPHelper(_clientSocket);
			else
				udphelper = new UDPHelper(new DatagramSocket());
			byte[] cookiebytes = sendCommandtoServer(udphelper, cmd, "NULL",
					null);
			String cookie = new String(cookiebytes);
			return sendCommandtoServer(udphelper, cmd, cookie, msg);
		} catch (Exception e) {
			e.printStackTrace();
			return null;
		}
	}

	private byte[] sendCommandtoServer(UDPHelper udphelper, String cmd,
			String cookie, byte[] msg) throws Exception {
		Packer packer = new Packer();
		packer.addElement(Packer.DataType.STRING, cmd);
		packer.addElement(Packer.DataType.STRING, cookie);
		if (msg != null)
			packer.addElement(Packer.DataType.BYTES, msg);
		SecretKey sk = null;
		if (!cmd.equalsIgnoreCase(LOGIN) && !cookie.equals("NULL"))
			sk = _serverSessionKey;
		udphelper.sendBytes(packer.pack(), _serverIp, _serverPort, sk);
		byte[] ret = udphelper.receiveBytes(20000, sk);
		_serverLoginIp = udphelper.fromip;
		_serverLoginPort = udphelper.fromport;
		return ret;
	}

	public ClientInfo searchOnlineByAddress(InetAddress ip) {
		try {
			// first update list
			updateList();

			for (ClientInfo c : _onlineUserList) {
				if (c._ip.equals(ip)) {
					return c;
				}
			}

			return null;
		} catch (Exception e) {
			e.printStackTrace();
			return null;
		}
	}

	public ClientInfo searchOnlineUserByName(String userName) {
		Iterator<ClientInfo> iter = _onlineUserList.iterator();
		while (iter.hasNext()) {
			ClientInfo c = iter.next();
			if (c._username.equals(userName))
				return c;
		}
		return null;
	}
}
